package com.ai.module.suno.service.impl;

import com.ai.module.chat.common.base.ThreadLocalContext;
import com.ai.module.chat.config.advice.exception.BusinessException;
import com.ai.module.chat.util.DateUtils;
import com.ai.module.chat.util.RedisTemplateUtils;
import com.ai.module.mj.constant.stringConstant.ReturnCode;
import com.ai.module.suno.config.SunoConfig;
import com.ai.module.suno.constant.SunoUrlConstant;
import com.ai.module.suno.domain.bo.SunoInfo;
import com.ai.module.suno.domain.bo.SunoInfoTask;
import com.ai.module.suno.domain.bo.TaskData;
import com.ai.module.suno.domain.entity.UserSunoEntity;
import com.ai.module.suno.domain.entity.UserSunoRecordEntity;
import com.ai.module.suno.domain.req.SunoReq;
import com.ai.module.suno.domain.vo.SunoGenerateVo;
import com.ai.module.suno.domain.vo.SunoRecordVo;
import com.ai.module.suno.domain.vo.UserSunoVo;
import com.ai.module.suno.mapper.UserSunoMapper;
import com.ai.module.suno.service.SunoRecordService;
import com.ai.module.suno.service.SunoService;
import com.alibaba.fastjson2.JSONObject;
import com.baomidou.mybatisplus.core.metadata.IPage;
import com.baomidou.mybatisplus.extension.service.impl.ServiceImpl;
import lombok.extern.slf4j.Slf4j;
import org.apache.tomcat.util.http.fileupload.IOUtils;
import org.jetbrains.annotations.NotNull;
import org.springframework.beans.BeanUtils;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.core.io.Resource;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.HttpClientErrorException;
import org.springframework.web.client.RestTemplate;

import javax.annotation.PostConstruct;
import javax.servlet.ServletOutputStream;
import javax.servlet.http.HttpServletResponse;
import java.io.InputStream;
import java.util.*;


/**
 * @author: heJinQin
 * @create: 2024-04-27 01:09
 * @description:
 **/
@Slf4j
@Service
public class SunoServiceImpl extends ServiceImpl<UserSunoMapper, UserSunoEntity> implements SunoService {

    public String TOKEN;
    private Map<String, String> COOKIE_MAP = new HashMap<>();

    @javax.annotation.Resource(name = "restTemplate")
    private RestTemplate restTemplate;

    @Autowired
    private SunoConfig sunoConfig;

//    @PostConstruct
    public void loadCookie() {
        Map<String, String> cookieMap = getRealCookie(sunoConfig.getCookie());
        // 查询 cookie
        loadCookieToken(cookieMap);
    }

    @NotNull
    private Map<String, String> getRealCookie(String cookie) {
        String[] cookieKV = cookie.split(";");
        for (String kv : cookieKV) {
            String[] kvArr = kv.split("=");
            if (kvArr.length < 2) continue;
            COOKIE_MAP.put(kvArr[0], kvArr[1]);
        }
        return COOKIE_MAP;
    }


    public void loadCookieToken(Map<String, String> cookieMap) {
        String url = SunoUrlConstant.refreshCookie + String.format(SunoUrlConstant.refreshToken, sunoConfig.getSessionId());
        ResponseEntity<JSONObject> response = postSuno(null, cookieMap, url);

        JSONObject body = response.getBody();
        log.info("模拟浏览器请求接口，返回数据：{}", body);
        HttpHeaders respHeader = response.getHeaders();
        List<String> respCookie = respHeader.get("Set-Cookie");
        String token = null;
        if (body != null) {
            token = body.getString("jwt");
            this.TOKEN = token;
            sunoConfig.setToken(token);
        }
        log.info("返回的 cookie [{}] token[{}]", respCookie, token);
        Map<String, String> realCookie1 = getRealCookie(String.join(";", respCookie));
        cookieMap.putAll(realCookie1);
        this.COOKIE_MAP = cookieMap;
    }

    @NotNull
    private ResponseEntity<JSONObject> postSuno(String param, Map<String, String> cookieMap, String url) {
        HttpHeaders headers = getSunoHeader(cookieMap);
        headers.setBearerAuth(sunoConfig.getToken());
//        headers.add(HttpHeaders.AUTHORIZATION,"Bearer "+sunoConfig.getToken());
        HttpEntity<String> entity = new HttpEntity<>(param, headers);
        ResponseEntity<JSONObject> response = restTemplate.postForEntity(url, entity, JSONObject.class);
        return response;
    }

    public static final String SUNO_RECORD = "suno:record:%s:%s";

    public static final String SUNO_Task = "suno:task:%s:%s";

    @Override
    public SunoInfoTask fetch(String sunoId) {
        try {
            String userId = ThreadLocalContext.getUserId();
            String sunoRecordKey = String.format(SUNO_RECORD, userId, sunoId);
            // 说明是有的
            if (RedisTemplateUtils.hasKey(sunoRecordKey)) {
                SunoInfoTask sunoInfoTask = RedisTemplateUtils.get(sunoRecordKey);

                boolean complete = checkNoSuccess(sunoInfoTask);
                if (!complete) {
                    // 说明是都成功了
                    return sunoInfoTask;
                }
                doDealFeedHttp(sunoId, sunoInfoTask, complete);
                return sunoInfoTask;
            }
            SunoInfoTask sunoInfoTask = new SunoInfoTask();

            String sunoDataKey = String.format(SUNO_Task, userId, sunoId);
            TaskData taskData = RedisTemplateUtils.get(sunoDataKey);
            BeanUtils.copyProperties(taskData, sunoInfoTask);
            doDealFeedHttp(sunoId, sunoInfoTask, false);
            return sunoInfoTask;
        } catch (HttpClientErrorException.Unauthorized exception) {
            // todo  刷新 token
            loadCookie();
            log.error("token 过期 重新加载 [{}], message [{}]", exception, exception.getMessage());
            return fetch(sunoId);
        }

    }

    private void doDealFeedHttp(String id, SunoInfoTask sunoInfoTask, boolean complete) {
        String userId = ThreadLocalContext.getUserId();
        String sunoRecordKey = String.format(SUNO_RECORD, userId, id);
        String generateUrl = SunoUrlConstant.fetchUrl + id;
        Map<String, String> cookieMap = getRealCookie(sunoConfig.getCookie());
        HttpHeaders headers = getSunoHeader(cookieMap);
        HttpEntity<Map<String, Object>> entity = new HttpEntity<>(headers);
        ResponseEntity<String> response = restTemplate.exchange(generateUrl, HttpMethod.GET, entity, String.class);
        String body = response.getBody();
        SunoInfo sunoInfoList = JSONObject.parseObject(body, SunoInfo.class);
        // 判断数据 是不是 成功 如果是 成功 则入库
        RedisTemplateUtils.set(sunoRecordKey, sunoInfoTask, DateUtils.TIME_DAY);
        String robotId = sunoInfoTask.getRobotId();
        String prompt = sunoInfoTask.getPrompt();
        BeanUtils.copyProperties(sunoInfoList, sunoInfoTask);
        sunoInfoTask.setPrompt(prompt);
        sunoInfoTask.setUserId(userId);
        sunoInfoTask.setRobotId(robotId);
        if (sunoInfoList != null) {
            complete = checkNoSuccess(sunoInfoTask);
        }
        if (!complete) {
            // 说明是都成功了
            List<UserSunoRecordEntity> userSunoRecordList = new ArrayList<>();
            // 进行数据入库 todo 这里改成 更新操作 会不会更好 ?
            UserSunoRecordEntity userSunoRecord = new UserSunoRecordEntity();
            BeanUtils.copyProperties(sunoInfoList, userSunoRecord);
            userSunoRecord.setSunoId(id);
            userSunoRecordList.add(userSunoRecord);
            // 进行数据入库
            sunoRecordService.saveBatch(userSunoRecordList);
        }
    }

    @Autowired
    private SunoRecordService sunoRecordService;

    private static boolean checkNoSuccess(SunoInfoTask sunoInfoTask) {
        return !sunoInfoTask.getStatus().equalsIgnoreCase("complete");
    }

    @Override
    public void download(String fileName, HttpServletResponse response) {
        try {
            String[] split = fileName.split("\\.");
            String fileType = split[split.length - 1];
            String URL = SunoUrlConstant.resourceUrl + fileName;
            HttpHeaders headers = new HttpHeaders();
//            headers.setContentType(MediaType.APPLICATION_JSON);
            HttpEntity request = new HttpEntity(headers);
            //当接受一个文件流时必须使用 resource (spring)接收
            ResponseEntity<Resource> entity = restTemplate.getForEntity(URL, Resource.class);
            InputStream in = entity.getBody().getInputStream();
            ServletOutputStream os = response.getOutputStream();
//            response.setHeader("Content-Type","video/mpeg");
//            response.setContentType("audio/mp3");
            if (fileType.equalsIgnoreCase("mp3")) {
                response.setContentType("video/mpeg");
            } else {
                response.setContentType("video/x-mpg");
            }
            response.setHeader(HttpHeaders.CONTENT_LENGTH, in.available() + "");
//            video/mpeg  video/mpeg4
            IOUtils.copy(in, os);
//            SunoService.downloadFile(in, fileName);
            in.close();
            os.flush();
            os.close();
        } catch (Exception exception) {

        }
    }

    @Override
    public JSONObject generate(SunoGenerateVo sunoGenerate) {
        String robotId = sunoGenerate.getRobotId();
        ResponseEntity<JSONObject> generatePost = postSuno(JSONObject.toJSONString(sunoGenerate), COOKIE_MAP, SunoUrlConstant.generateUrl);
        HttpStatus statusCode = generatePost.getStatusCode();
        if (statusCode != HttpStatus.OK) {
            log.error("任务失败");
            throw new BusinessException("AI音乐处理失败", ReturnCode.FAILURE);
        }
        JSONObject generateJson = generatePost.getBody();
        log.info("suno-ai结果为 [{}]", Objects.requireNonNull(generateJson).toString());
        //{"id":"f5763e63-698e-48b0-b16d-99b9b15ab7c6","clips":[{"id":"768590b2-6afe-4c66-8461-680d78cbb5a1","video_url":"","audio_url":"","is_video_pending":false,"major_model_version":"v3","model_name":"chirp-v3","metadata":{"tags":"mellow hip hop","prompt":"唱一首歌 今天你要嫁给我呀","type":"gen","stream":true},"is_liked":false,"user_id":"9cf1cf7d-cc58-427f-a12e-72d826ae7453","display_name":"RedHotRiff706","handle":"redhotriff706","is_handle_updated":false,"is_trashed":false,"created_at":"2024-04-26T18:04:25.104Z","status":"submitted","title":"test-singel","play_count":0,"upvote_count":0,"is_public":false},{"id":"132fb793-64da-4cbc-b9bb-993a3c382dae","video_url":"","audio_url":"","is_video_pending":false,"major_model_version":"v3","model_name":"chirp-v3","metadata":{"tags":"mellow hip hop","prompt":"唱一首歌 今天你要嫁给我呀","type":"gen","stream":true},"is_liked":false,"user_id":"9cf1cf7d-cc58-427f-a12e-72d826ae7453","display_name":"RedHotRiff706","handle":"redhotriff706","is_handle_updated":false,"is_trashed":false,"created_at":"2024-04-26T18:04:25.104Z","status":"submitted","title":"test-singel","play_count":0,"upvote_count":0,"is_public":false}],"metadata":{"tags":"mellow hip hop","prompt":"唱一首歌 今天你要嫁给我呀","type":"gen","stream":true},"major_model_version":"v3","status":"running","created_at":"2024-04-26T18:04:25.090Z","batch_size":2}
        // 生成记录 入库
        String realId = generateJson.getString("id");
        List<JSONObject> clips = generateJson.getJSONArray("clips").toJavaList(JSONObject.class);
        clips.forEach(clip -> {
            String sunoId = clip.getString("id");
            TaskData taskData = new TaskData();
            taskData.setUserId(ThreadLocalContext.getUserId());
            taskData.setRobotId("test");
            taskData.setSunoId(sunoId);
            String userId = ThreadLocalContext.getUserId();
            String sunoRecordKey = String.format(SUNO_Task, userId, sunoId);
            RedisTemplateUtils.set(sunoRecordKey, taskData, DateUtils.TIME_DAY);
        });
        // 生成 suno
        JSONObject jsonObject = clips.get(0);
        UserSunoEntity userSuno = new UserSunoEntity();
        userSuno.setTaskId(realId);
        userSuno.setModel(jsonObject.getString("model_name"));
        userSuno.setPrompt(jsonObject.getJSONObject("metadata").getString("prompt"));
        userSuno.setTags(jsonObject.getJSONObject("metadata").getString("tags"));
        userSuno.setTitle(jsonObject.getString("title"));
        userSuno.setUserId(ThreadLocalContext.getUserId());
        userSuno.setRobotId(robotId);
        save(userSuno);
        return generateJson;
    }

    public HttpHeaders getSunoHeader(Map<String, String> cookieMap) {
        List<String> cookieList = new ArrayList<>();
        for (String key : cookieMap.keySet()) {
            String cookie = String.format("%s=%s", key, cookieMap.get(key));
            cookieList.add(cookie);
        }
        String realCookie = String.join(";", cookieList);
        HttpHeaders headers = new HttpHeaders();
        List<String> cookies = new ArrayList<>();
        // 登录获取Cookie 这里是直接给Cookie，可使用下方的login方法拿到Cookie给入
        // 把上一步获取的cookie，在请求header里面传入
        cookies.add(realCookie);
        // 将cookie存入头部
        headers.put(HttpHeaders.COOKIE, cookies);
        headers.setContentType(MediaType.APPLICATION_JSON);
        headers.setBearerAuth(sunoConfig.getToken());
//        headers.add(HttpHeaders.AUTHORIZATION,"Bearer "+ sunoConfig.getToken());
        return headers;
    }

    @Override
    public IPage<UserSunoVo> sunoRecord(SunoReq page) {
        return sunoRecordService.sunoRecord(page);
    }

    @Override
    public SunoRecordVo sunoDetail(String sunoId) {
        return sunoRecordService.sunoDetail(sunoId);
    }
}
